package com.liaoyifan.core.thread;

import java.util.*;
import java.util.concurrent.*;
import java.util.concurrent.atomic.AtomicReference;
import java.util.function.Consumer;
import java.util.function.Function;
import lombok.extern.log4j.Log4j2;

@SuppressWarnings("unused")
@Log4j2
public class ThreadPool {

    private static final Map<String, CorePool> THREAD_POOL_CACHE = new ConcurrentHashMap<>();

    // 注册JVM关闭钩子
    static {
        Runtime.getRuntime().addShutdownHook(new Thread(ThreadPool::shutdownAll));
    }

    // 获取或创建线程池
    public static ExecutorService get(PoolConfig config) {
        return THREAD_POOL_CACHE
                .computeIfAbsent(config.getPoolName(), name -> CorePool.create(config))
                .getExecutor();
    }

    // 执行任务
    public static void execute(String poolName, Consumer<ExecutorService> consumer) {
        execute(PoolConfig.get(poolName), consumer);
    }

    public static void execute(PoolConfig config, Consumer<ExecutorService> consumer) {
        ExecutorService executor = get(config);
        try {
            consumer.accept(executor);
        } finally {
            if (config.isShutdown()) {
                shutdown(config.getPoolName());
            }
        }
    }

    // 提交任务
    public static <R> R submit(String poolName, Function<ExecutorService, R> func) {
        return submit(PoolConfig.get(poolName), func);
    }

    public static <R> R submit(PoolConfig config, Function<ExecutorService, R> func) {
        ExecutorService executor = get(config);
        try {
            return func.apply(executor);
        } finally {
            if (config.isShutdown()) {
                shutdown(config.getPoolName());
            }
        }
    }

    public static <T> void forEach(String poolName, List<T> list, Consumer<T> consumer) {
        forEach(PoolConfig.get(poolName), list, consumer);
    }

    public static <T> void forEach(PoolConfig config, List<T> list, Consumer<T> consumer) {
        if (list == null || list.isEmpty() || consumer == null) {
            return;
        }
        ExecutorService executor = get(config.corePoolSize(list.size()));
        try {
            internalForEach(executor, list, consumer);
        } finally {
            if (config.isShutdown()) {
                shutdown(config.getPoolName());
            }
        }
    }

    private static <T> void internalForEach(
            ExecutorService executor, List<T> list, Consumer<T> consumer) {
        AtomicReference<Exception> firstError = new AtomicReference<>();
        List<CompletableFuture<Void>> futures =
                list.stream()
                        .map(
                                item ->
                                        CompletableFuture.runAsync(
                                                () -> {
                                                    try {
                                                        if (firstError.get() != null) {
                                                            return; // 如果已发生错误，跳过后续处理
                                                        }
                                                        consumer.accept(item);
                                                    } catch (Exception e) {
                                                        firstError.compareAndSet(null, e);
                                                        throw new CompletionException(e);
                                                    }
                                                },
                                                executor))
                        .toList();

        CompletableFuture.allOf(futures.toArray(new CompletableFuture[0]))
                .exceptionally(
                        ex -> {
                            Throwable cause = ex.getCause() != null ? ex.getCause() : ex;
                            firstError.compareAndSet(null, new RuntimeException(cause));
                            return null;
                        })
                .join();

        if (firstError.get() != null) {
            throw new RuntimeException("并行处理过程中发生异常", firstError.get());
        }
    }

    public static void shutdown(String poolName) {
        PoolConfig.remove(poolName);
        CorePool pool = THREAD_POOL_CACHE.remove(poolName);
        if (pool != null) {
            pool.shutdown();
        }
    }

    public static void shutdownAll() {
        log.info("Shutting down all thread pools...");
        PoolConfig.clear();
        List<CorePool> pools = new ArrayList<>(THREAD_POOL_CACHE.values());
        THREAD_POOL_CACHE.clear();
        pools.forEach(CorePool::shutdown);
    }
}
